#Author: CHEN Xin

import matplotlib.pyplot as plt
from pylab import *
from matplotlib.finance import candlestick_ochl
from matplotlib import gridspec,patches
import fetchdata

def DrawBars(ax,poslist,valuelist,up="salmon",down="darkseagreen"):
    width=0.7
    for i in xrange(len(poslist)):
        pos=poslist[i]
        value=valuelist[i]
        color=up
        if value<0:
            color=down
        patch = ax.add_patch(
            patches.Rectangle(
                (pos - width / 2, 0),  # (x,y)
                width,  # width
                value,  # height
                edgecolor=color,
                facecolor=color
            ),
        )

def DrawBOLL(ax,datas):
    poslist=[one["pos"] for one in datas]
    upband=[one["bband"][0] for one in datas]
    ave=[one["bband"][1] for one in datas]
    dnband=[one["bband"][2] for one in datas]
    ax.plot(poslist,ave,color="b",alpha=0.5)
    ax.plot(poslist,upband,color='r',alpha=0.5)
    ax.plot(poslist,dnband,color="g",alpha=0.5)
def DrawMacd(ax,datas):
    poslist=[one["pos"] for one in datas]
    macd=[one["macd"][0] for one in datas]
    macdsignal=[one["macd"][1] for one in datas]
    macdhist=[one["macd"][2]*2 for one in datas]
    DrawBars(ax, poslist, macdhist)
    ax.plot(poslist,macd,color="r")
    ax.plot(poslist, macdsignal, color="g")
def DrawSAR(ax,datas):
    poslist=[one["pos"] for one in datas]
    psar=[one["sar"] for one in datas]
    ax.plot(poslist, psar,  "o", markersize=3, markeredgewidth=0, markeredgecolor='r', markerfacecolor='r')
def DrawSTOCH(ax,datas):
    poslist=[one["pos"] for one in datas]
    slowk=[one["stoch"][0] for one in datas]
    slowd=[one["stoch"][1] for one in datas]
    ax.plot(poslist,slowk,color="r")
    ax.plot(poslist, slowd, color="g")
def DrawTranPoint(ax,datas):
    for one in datas:
        if "tran_mark" in one:
            if one["tran_mark"]==1:
                ax.annotate("%0.3f"%one["tran_mark"], xy=(one["pos"] , one["Low"]-5), xytext=(one["pos"] , one["Low"]-7),
                            arrowprops=dict(facecolor='red', width=1, shrink=0.25))
            if one["tran_mark"]==-1:
                ax.annotate("%0.3f"%one["tran_mark"], xy=(one["pos"] , one["High"]+5), xytext=(one["pos"] , one["High"]+7),
                            arrowprops=dict(facecolor='green', width=1, shrink=0.25))
def exDrawLine(axs,spandata):
    DrawBOLL(axs[0],spandata)
    DrawSAR(axs[0],spandata)
    DrawMacd(axs[2],spandata)
    DrawSTOCH(axs[1],spandata)
    DrawTranPoint(axs[0],spandata)
def Show(datalist,exDrawLine=None,span=100):
    spandata=datalist[-span:]
    datas=[]
    for i in xrange(len(spandata)):
        spandata[i]["pos"]=i
        datas.append([i,spandata[i]["Open"],spandata[i]["Close"],spandata[i]["High"],spandata[i]["Low"],spandata[i]["Date"]])
    idtodate={}
    idtodate = {one[0]: one[5] for one in datas}

    fig=plt.figure()
    fig.subplots_adjust(top=0.98,bottom=0.02,left=0.05,right=0.98)
    fig.subplots_adjust(hspace=0.01)
    gs = gridspec.GridSpec(3, 1, height_ratios=[4, 1,1])
    ax = plt.subplot(gs[0])
    ax.set_ylabel("USD/oz")
    ax1 = plt.subplot(gs[1], sharex=ax)
    ax2 = plt.subplot(gs[2], sharex=ax)
    def getDate(x,position):
        if x in idtodate:
            return idtodate[x]
        return ""
    ax.xaxis.set_major_formatter(FuncFormatter(getDate))
    bars=candlestick_ochl(ax, datas, width=0.6,colorup='r', colordown='g')

    if exDrawLine:
        exDrawLine((ax,ax1,ax2),spandata)

    ax.autoscale_view()
    ax1.autoscale_view()
    ax2.autoscale_view()

    ax.grid()
    ax1.grid()
    ax2.grid()

    mng = plt.get_current_fig_manager()
    mng.full_screen_toggle()
    plt.show()

if __name__ == "__main__":
    datalist=fetchdata.loaddate("data/golddata.gz")
    Show(datalist,exDrawLine,500)
    Show("data/goldmin.svr", "data/goldmin.data", True, 1200)
